#include "AsyncIOManager.h"

AsyncIOManager::AsyncIOManager(size_t threadCount) : m_shutdown(false) {
    // 创建工作线程
    for (size_t i = 0; i < threadCount; ++i) {
        m_workers.emplace_back(&AsyncIOManager::WorkerThread, this);
    }
}

AsyncIOManager::~AsyncIOManager() {
    {
        std::lock_guard<std::mutex> lock(m_mutex);
        m_shutdown = true;
    }
    m_condition.notify_all();
    
    // 等待所有工作线程结束
    for (auto& worker : m_workers) {
        if (worker.joinable()) {
            worker.join();
        }
    }
}

void AsyncIOManager::QueueRead(const std::wstring& path, std::vector<uint8_t>& buffer,
                             size_t offset, size_t length,
                             std::function<void(NTSTATUS)> callback) {
    IORequest req;
    req.type = IORequest::READ;
    req.path = path;
    req.buffer = std::move(buffer);
    req.offset = offset;
    req.length = length;
    req.callback = std::move(callback);

    {
        std::lock_guard<std::mutex> lock(m_mutex);
        m_pendingRequests.push(std::move(req));
    }
    m_condition.notify_one();
}

void AsyncIOManager::QueueWrite(const std::wstring& path, const std::vector<uint8_t>& buffer,
                              size_t offset, std::function<void(NTSTATUS)> callback) {
    IORequest req;
    req.type = IORequest::WRITE;
    req.path = path;
    req.buffer = buffer;
    req.offset = offset;
    req.callback = std::move(callback);

    {
        std::lock_guard<std::mutex> lock(m_mutex);
        m_pendingRequests.push(std::move(req));
    }
    m_condition.notify_one();
}

void AsyncIOManager::WorkerThread() {
    while (true) {
        IORequest req;
        
        {
            std::unique_lock<std::mutex> lock(m_mutex);
            m_condition.wait(lock, [this] {
                return m_shutdown || !m_pendingRequests.empty();
            });
            
            if (m_shutdown && m_pendingRequests.empty()) {
                return;
            }
            
            req = std::move(m_pendingRequests.front());
            m_pendingRequests.pop();
        }
        
        ProcessRequest(req);
    }
}

void AsyncIOManager::ProcessRequest(IORequest& req) {
    HANDLE hFile = CreateFileW(
        req.path.c_str(),
        req.type == IORequest::READ ? GENERIC_READ : GENERIC_WRITE,
        FILE_SHARE_READ | FILE_SHARE_WRITE,
        nullptr,
        OPEN_EXISTING,
        FILE_FLAG_OVERLAPPED,
        nullptr
    );

    if (hFile == INVALID_HANDLE_VALUE) {
        req.callback(STATUS_UNSUCCESSFUL);
        return;
    }

    OVERLAPPED overlapped = {0};
    overlapped.Offset = static_cast<DWORD>(req.offset & 0xFFFFFFFF);
    overlapped.OffsetHigh = static_cast<DWORD>(req.offset >> 32);

    BOOL success;
    if (req.type == IORequest::READ) {
        success = ReadFile(
            hFile,
            req.buffer.data(),
            static_cast<DWORD>(req.length),
            nullptr,
            &overlapped
        );
    } else {
        success = WriteFile(
            hFile,
            req.buffer.data(),
            static_cast<DWORD>(req.buffer.size()),
            nullptr,
            &overlapped
        );
    }

    if (!success && GetLastError() != ERROR_IO_PENDING) {
        CloseHandle(hFile);
        req.callback(STATUS_UNSUCCESSFUL);
        return;
    }

    // 等待IO完成
    DWORD bytesTransferred;
    success = GetOverlappedResult(hFile, &overlapped, &bytesTransferred, TRUE);
    
    CloseHandle(hFile);
    req.callback(success ? STATUS_SUCCESS : STATUS_UNSUCCESSFUL);
}
